import os
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default='0', help='CUDA device(s) to use')
args, _ = parser.parse_known_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import  VLChatProcessor
from PIL import Image

import cv2
import numpy as np
osj = os.path.join
ose = os.path.exists


def janus_pro_7b():
    # Load model and processor
    # model_path = "deepseek-ai/Janus-Pro-7B"
    model_path = "/mnt/nas/zhangshu/Janus-Pro-7B"
    config = AutoConfig.from_pretrained(model_path)
    language_config = config.language_config
    language_config._attn_implementation = 'eager'
    vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
                                                language_config=language_config,
                                                trust_remote_code=True)
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()

    vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
    tokenizer = vl_chat_processor.tokenizer

    return vl_chat_processor, tokenizer, vl_gpt


@torch.inference_mode()
def multimodal_understanding(vl_chat_processor, tokenizer, vl_gpt,
    image_path = '',
    image=None, question='', seed=42, top_p=0.95, temperature=0.1,
    max_new_tokens = 512):

    image = image if image is not None else np.array( Image.open(image_path) )

    cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'
    # Clear CUDA cache before generating
    torch.cuda.empty_cache()
    
    # set seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)
    
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]
    
    pil_images = [Image.fromarray(image).convert('RGB')]
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
    
    
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
    
    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=max_new_tokens,
        do_sample=False if temperature == 0 else True,
        use_cache=True,
        temperature=temperature,
        top_p=top_p,
    )
    
    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return answer


if __name__=='__main__':
    import argparse,os
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--cuda', type=str, default='2', help='CUDA device id')
    args, unknown = parser.parse_known_args()
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda

    vl_chat_processor, tokenizer, vl_gpt = janus_pro_7b()

    ans = multimodal_understanding(
        vl_chat_processor, tokenizer, vl_gpt,
        image_path = 'compare_pockets/np-3.jpg',
        question='描述这件衣服的颜色'
    )
    print(ans)
